#include "execution/plans/limit_plan.h"
#include "execution/plans/sort_plan.h"
#include "execution/plans/topn_plan.h"
#include "optimizer/optimizer.h"

namespace bustub {

auto Optimizer::OptimizeSortLimitAsTopN(const AbstractPlanNodeRef &plan) -> AbstractPlanNodeRef {
  std::vector<AbstractPlanNodeRef> children;
  for (const auto &child : plan->GetChildren()) {
    children.emplace_back(OptimizeSortLimitAsTopN(child));
  }
  auto optimized_plan = plan->CloneWithChildren(std::move(children));
  auto limit_plan = dynamic_cast<LimitPlanNode *>(optimized_plan.get());

  if (optimized_plan->GetType() == PlanType::Limit && limit_plan->GetChildPlan()->GetType() == PlanType::Sort) {
    auto sort_plan = dynamic_cast<const SortPlanNode *>(limit_plan->GetChildPlan().get());
    LOG_DEBUG("========");
    return std::make_shared<TopNPlanNode>(std::make_shared<Schema>(limit_plan->OutputSchema()),
                                          sort_plan->GetChildPlan(), sort_plan->GetOrderBy(), limit_plan->GetLimit());
  }
  return optimized_plan;
}

}  // namespace bustub
